from sklearn.neural_network import MLPRegressor
from sklearn.datasets import make_regression
from sklearn.model_selection import train_test_split
import pandas as pd
import numpy as np


def get_dataset():
    humidity_depth = '10'  # '10' '40', '100', '200'
    humidity_col = ['year', humidity_depth]
    humidity_raw = pd.read_excel('data/humidity.xls', names=['year', 'lon', 'lat', '10', '40', '100', '200'])
    humidity_raw = humidity_raw[humidity_col]
    steam_col = ['year', 'wm2', 'steam_mm']
    steam_raw = pd.read_excel('data/steam.xls', names=['mon', 'year', 'lon', 'lat', 'wm2', 'steam_mm'])
    weather_use_col = ['降水量(mm)', '降水天数']
    weather_col = ['raw_mm', 'rain_day']

    weather = pd.DataFrame(columns=weather_col)
    steam = pd.DataFrame(columns=steam_col)
    humidity = pd.DataFrame(columns=humidity_col)
    years = [2012, 2013, 2014, 2015, 2016, 2017, 2018, 2019, 2020, 2021, 2022]
    for year in years:
        humidity = humidity.append(humidity_raw[humidity_raw.year == year], ignore_index=True)
        steam = steam.append(steam_raw[steam_raw.year == year], ignore_index=True)
        weather_raw = pd.read_excel(f'data/weather/{year}年.xls', usecols=weather_use_col)
        weather_raw.columns = weather_col
        weather = weather.append(weather_raw, ignore_index=True)
    # 数据整合
    for col in steam_col:
        humidity[col] = steam[col]
    for col in weather_col:
        humidity[col] = weather[col]

    target = humidity[humidity_depth]
    return np.array(humidity.to_numpy(), dtype=float), np.array(target.to_numpy(), dtype=float)


X, y = get_dataset()
X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=1)
regr = MLPRegressor(random_state=1, max_iter=1000).fit(X_train, y_train)
prediction = regr.predict(X_test[:10])
print(prediction)
print(y_test[:10])
